from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.metrics import roc_auc_score, classification_report
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.model_selection import GridSearchCV
import joblib
from xgboost import XGBClassifier
from sklearn.ensemble import RandomForestClassifier

plt.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体为黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像时负号 '-' 显示为方块的问题

# 读取数据
data = pd.read_csv("../../../data/raw/train.csv")
# 数据预处理  将有多分类的列 转化为 数值型的
label_encoder = LabelEncoder()
data = data.apply(label_encoder.fit_transform)

# 选取特征列  22列
X = data[["BusinessTravel",
          "Department",
          "EducationField",
          "EnvironmentSatisfaction",
          "JobInvolvement",
          "JobLevel",
          "JobRole",
          "JobSatisfaction",
          "MaritalStatus",
          "OverTime",
          "StockOptionLevel",
          "WorkLifeBalance",

          "Age",
          "DistanceFromHome",
          "MonthlyIncome",
          "NumCompaniesWorked",
          "PercentSalaryHike",
          "TotalWorkingYears",
          # "YearsAtCompany",
          "YearsInCurrentRole",
          "YearsSinceLastPromotion",
          "YearsWithCurrManager"
          ]
]
y = data["Attrition"]
# 对X 标准化
scale = StandardScaler()
scale.fit(X)
X = scale.transform(X)
joblib.dump(scale, "../model/scale_model.pkl")


# 划分数据集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=25)
# 使用XGBClassifier 进行模型训练
# es_xgb = XGBClassifier(
#     # learning_rate=0.1,
#     random_state=3,
#     n_estimators=300,
#     max_depth=1,  # 2
#     min_child_weight=2,
#     gamma=0,
#     reg_alpha=0,
#     reg_lambda=1.0,
#     objective='binary:logistic',
# )
es_xgb = RandomForestClassifier(n_estimators=500,
                                    min_samples_split=5,
                                    min_samples_leaf=1,
                                    max_features='log2',
                                    max_depth=20,
                                    random_state=22,
                                    n_jobs=-1)

es_xgb.fit(X_train, y_train)
y_pred = es_xgb.predict_proba(X_test)[:, 1]

# 打印分类评估报告
# print("分类评估报告：")
# print(classification_report(y_test, y_pred))
print("AUC:", roc_auc_score(y_test, y_pred))  # AUC: 0.899702229646307

# 保存模型
joblib.dump(es_xgb, "xgb_model.pkl")
